// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Authors: Alessandro Tasora
// =============================================================================

#ifndef CHGNUPLOT_H
#define CHGNUPLOT_H

#include <sstream>
#include <iostream>

#include "chrono/core/ChStream.h"
#include "chrono/core/ChMatrixDynamic.h"
#include "chrono/core/ChVectorDynamic.h"
#include "chrono/assets/ChColor.h"
#include "chrono/motion_functions/ChFunction_Base.h"
#include "chrono/motion_functions/ChFunction_Recorder.h"
#include "chrono/motion_functions/ChFunction_Oscilloscope.h"

#include "chrono_postprocess/ChApiPostProcess.h"

namespace chrono {

/// Namespace with classes for the postprocess unit.
namespace postprocess {

/// @addtogroup postprocess_module
/// @{

class ChGnuPlotDataplot {
  public:
    std::string command;
    ChMatrixDynamic<double> data;
};

/// Class for lotting data with GNUplot.
/// This is a basic utility class, it just saves a temporary gpl file on disk and then
/// calls the GNUplot utility from the system shell, so it does not require to link
/// GNUplot libraries. If the GNUplot is not installed, simply nothing happens.
/// NOTE: the GNUplot should be installed in your platform. Supported from release 4.6.
/// NOTE: the GNUplot executable must be available in your PATH, i.e. if you type 'gnuplot' in the
/// cmd window, it should be executed. If this does not happens, add the directory with the GNUplot
/// binary to your PATH. In Linux platforms this should be already ok by default.
/// NOTE: if no plot is displayed, open a cmd shell, type "gnuplot __tmp_gnuplot.gpl -persist"
///  and see which error is displayed.

class ChGnuPlot {
  public:
    ChGnuPlot(const char* mgpl_filename = "__tmp_gnuplot.gpl") {
        commandfile += "# This is an autogenerated .gpl file that is executed by GNUplot \n";
        commandfile += "# It is created by the ChGnuPlot.h class of Chrono::Engine \n\n";
        persist = true;
        gpl_filename = mgpl_filename;
    }

    virtual ~ChGnuPlot() {
        FlushPlots(commandfile);
        ExecuteGnuplot(this->commandfile);
    }

    /// If false, the GNUplot will close soon after the plot. For plotting in
    /// windows , it is better to leave it as true by default, but for plotting eps or png
    /// maybe you want it to close automatically.
    void SetPersist(bool mpersist) { this->persist = mpersist; }

    /// This is the main feature.
    /// Use this function to add commands in the gnuplot script.
    /// It can be whatever gnuplot statement.
    /// Basically you would just need calls to SetCommand() followed by an Output()
    /// at the end, however to make things easier, there are some shortcut functions
    /// such as SetRangeX() ,SetGrid(), Plot etc. that populate the command file for you.
    /// It will be automatically null-terminated, so you may call it multiple times,
    /// each per line.
    void SetCommand(const char* command) {
        commandfile += command;
        commandfile += "\n";
    }

    /// This is equivalent to SetCommand().
    ChGnuPlot& operator<<(const char* command) {
        this->SetCommand(command);
        return *this;
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from a .dat external file
    void Plot(const char* datfile, int colX, int colY, const char* title, const char* customsettings = " with lines ") {
        ChGnuPlotDataplot mdataplot;
        mdataplot.data.Resize(0, 0);  // embedded data not needed

        mdataplot.command += " \"";
        mdataplot.command += datfile;
        mdataplot.command += "\" using ";
        mdataplot.command += i_to_str(colX);
        mdataplot.command += ":";
        mdataplot.command += i_to_str(colY);
        mdataplot.command += " ";
        mdataplot.command += customsettings;
        mdataplot.command += " title \"";
        mdataplot.command += title;
        mdataplot.command += "\" ";

        this->plots.push_back(mdataplot);
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from two column vectors
    void Plot(ChVectorDynamic<>& mx,
              ChVectorDynamic<>& my,
              const char* title,
              const char* customsettings = " with lines ") {
        assert(mx.GetRows() == my.GetRows());

        ChGnuPlotDataplot mdataplot;
        mdataplot.data.Resize(mx.GetRows(), 2);
        mdataplot.data.PasteMatrix(mx, 0, 0);
        mdataplot.data.PasteMatrix(my, 0, 1);

        mdataplot.command += " \"-\" using 1:2 ";
        mdataplot.command += customsettings;
        mdataplot.command += " title \"";
        mdataplot.command += title;
        mdataplot.command += "\" ";

        this->plots.push_back(mdataplot);
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from two columns of a matrix
    void Plot(ChMatrix<>& mdata, int colX, int colY, const char* title, const char* customsettings = " with lines ") {
        ChVectorDynamic<> mx(mdata.GetRows());
        ChVectorDynamic<> my(mdata.GetRows());
        mx.PasteClippedMatrix(mdata, 0, colX, mdata.GetRows(), 1, 0, 0);
        my.PasteClippedMatrix(mdata, 0, colY, mdata.GetRows(), 1, 0, 0);
        Plot(mx, my, title, customsettings);
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from a ChFunction_recorder
    void Plot(ChFunction_Recorder& mrecorder, const char* title, const char* customsettings = " with lines ") {
        ChVectorDynamic<> mx((const int)mrecorder.GetPoints().size());
        ChVectorDynamic<> my(mx.GetRows());

        int i = 0;
        for (auto iter = mrecorder.GetPoints().begin(); iter != mrecorder.GetPoints().end(); ++iter) {
            mx(i) = iter->x;
            my(i) = iter->y;
            ++i;
        }
        Plot(mx, my, title, customsettings);
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from a ChFunction_recorder
    void Plot(ChFunction_Oscilloscope& mrecorder, const char* title, const char* customsettings = " with lines ") {
        ChVectorDynamic<> mx((int)(mrecorder.GetPointList().size()));
        ChVectorDynamic<> my(mx.GetRows());

        double xmin, xmax;
        mrecorder.Estimate_x_range(xmin, xmax);
        double x = xmin;
        int i = 0;
        std::list<double>::iterator iter = mrecorder.GetPointList().begin();
        while (iter != mrecorder.GetPointList().end()) {
            mx(i) = x;
            my(i) = (*iter);
            x += mrecorder.Get_dx();
            ++iter;
            ++i;
        }
        Plot(mx, my, title, customsettings);
    }

    /// Shortcut to easy 2D plot of x,y data
    /// from a generic ChFunction
    /// Note, if the ChFunction is of type ChFunction_Oscilloscope or ChFunction_Recorder, there
    /// are specific Plot() functions that can leverage their point-like nature in a better way.
    void Plot(ChFunction& mfunct,
              double xmin,
              double xmax,
              double dx,
              const char* title,
              const char* customsettings = " with lines ") {
        int samples = (int)floor((xmax - xmin) / dx);
        ChVectorDynamic<> mx(samples);
        ChVectorDynamic<> my(samples);

        double x = xmin;
        for (int i = 0; i < samples; ++i) {
            mx(i) = x;
            my(i) = mfunct.Get_y(x);
            x += dx;
        }
        Plot(mx, my, title, customsettings);
    }

    /// Replot last plots (since the last time you used Output...() functions, if any).
    void Replot() { commandfile += "replot \n"; }

    /// Shortcut to add the command that set the X range
    void SetRangeX(double mmin, double mmax, bool automin = false, bool automax = false) {
        commandfile += "set xrange [";
        if (automin)
            commandfile += "*";
        else
            commandfile += d_to_str(mmin);
        commandfile += ":";
        if (automax)
            commandfile += "*";
        else
            commandfile += d_to_str(mmax);
        commandfile += "] \n";
    }
    /// Shortcut to add the command that set the Y range
    void SetRangeY(double mmin, double mmax, bool automin = false, bool automax = false) {
        commandfile += "set yrange [";
        if (automin)
            commandfile += "*";
        else
            commandfile += d_to_str(mmin);
        commandfile += ":";
        if (automax)
            commandfile += "*";
        else
            commandfile += d_to_str(mmax);
        commandfile += "] \n";
    }
    /// Shortcut to add the command that set the Z range
    void SetRangeZ(double mmin, double mmax, bool automin = false, bool automax = false) {
        commandfile += "set zrange [";
        if (automin)
            commandfile += "*";
        else
            commandfile += d_to_str(mmin);
        commandfile += ":";
        if (automax)
            commandfile += "*";
        else
            commandfile += d_to_str(mmax);
        commandfile += "] \n";
    }

    /// Shortcut to add title
    void SetTitle(const char* mlabel) {
        commandfile += "set title \"";
        commandfile += mlabel;
        commandfile += "\" \n";
    }

    /// Shortcut to add a label on axis
    void SetLabelX(const char* mlabel) {
        commandfile += "set xlabel \"";
        commandfile += mlabel;
        commandfile += "\" \n";
    }

    /// Shortcut to add a label on axis
    void SetLabelY(const char* mlabel) {
        commandfile += "set ylabel \"";
        commandfile += mlabel;
        commandfile += "\" \n";
    }

    /// Shortcut to add a label on axis
    void SetLabelZ(const char* mlabel) {
        commandfile += "set zlabel \"";
        commandfile += mlabel;
        commandfile += "\" \n";
    }

    /// Shortcut to add the command that turns on the grid
    void SetGrid(bool dashed = true, double linewidth = 1.0, ChColor mcolor = ChColor(0, 0, 0)) {
        commandfile += "set grid ";
        if (dashed)
            commandfile += "lt 0 ";
        else
            commandfile += "lt 1 ";
        commandfile += "lw ";
        commandfile += d_to_str(linewidth);
        commandfile += " lc ";
        commandfile += col_to_hex(mcolor);
        commandfile += "\n";
    }

    /// Set plot in a window.
    /// For multiple windows, call this with increasing windownum, interleaving with Plot() statements etc.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputWindow(int windownum = 0) {
        FlushPlots(commandfile);
        commandfile += "set terminal wxt ";
        commandfile += i_to_str(windownum);
        commandfile += "\n";
    }

    /// Save plot in a png file.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputPNG(const char* filename, int sizex = 400, int sizey = 300) {
        FlushPlots(commandfile);
        commandfile += "set terminal png size ";
        commandfile += sizex;
        commandfile += ",";
        commandfile += sizey;
        commandfile += "\n";
        commandfile += "set output '";
        commandfile += filename;
        commandfile += "'\n";
    }

    /// Save plot in a eps file.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputEPS(const char* filename, double inchsizex = 4, double inchsizey = 3, bool color = true) {
        FlushPlots(commandfile);
        commandfile += "set terminal postscript eps ";
        if (color)
            commandfile += " color ";
        commandfile += " size ";
        commandfile += d_to_str(inchsizex);
        commandfile += ",";
        commandfile += d_to_str(inchsizey);
        commandfile += "\n";
        commandfile += "set output '";
        commandfile += filename;
        commandfile += "'\n";
    }

    /// Save plot in a custom terminal.
    /// For instance try terminalsetting ="set terminal svg size 350,262 fname 'Verdana' fsize 10"
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputCustomTerminal(const char* filename, const char* terminalsetting) {
        FlushPlots(commandfile);
        commandfile += terminalsetting;
        commandfile += "set output '";
        commandfile += filename;
        commandfile += "'\n";
    }

  protected:
    void FlushPlots(std::string& mscript) {
        // generate the  plot xxx , yyy , zzz stuff:
        if (plots.size() > 0)
            mscript += "plot ";
        for (int i = 0; i < this->plots.size(); ++i) {
            if (i > 0)
                mscript += " ,\\\n";
            mscript += this->plots[i].command;
        }
        mscript += "\n";

        // Embed plot data in the .gpl file
        for (int i = 0; i < this->plots.size(); ++i) {
            if ((plots[i].data.GetColumns() > 0) && (plots[i].data.GetRows() > 0)) {
                for (int ir = 0; ir < plots[i].data.GetRows(); ++ir) {
                    for (int ic = 0; ic < plots[i].data.GetColumns(); ++ic) {
                        mscript += d_to_str(plots[i].data(ir, ic));
                        mscript += " ";
                    }
                    mscript += "\n";
                }
                mscript += "end\n";
            }
        }
        this->plots.resize(0);
    }

    void ExecuteGnuplot(std::string& mscript) {
        // Create a tmp .gpl file
        {
            ChStreamOutAsciiFile gnuplot_command(this->gpl_filename.c_str());
            gnuplot_command << mscript;
        }

        std::string syscmd;

// Launch the GNUplot from shell:
#ifdef _WIN32
        // ex. of launched sys command: "start gnuplot __tmp_gnuplot.gpl -persist"
        syscmd += "start gnuplot \"";
        syscmd += this->gpl_filename;
        syscmd += "\"";
        if (persist)
            syscmd += " -persist";
        system(syscmd.c_str());
#else
        // Unix like systems:
        // ex. of launched sys command: "gnuplot __tmp_gnuplot.gpl -persist &"
        syscmd += "gnuplot \"";
        syscmd += this->gpl_filename;
        syscmd += "\"";
        if (persist)
            syscmd += " -persist";
        syscmd += " &";  // to launch and forget
        system(syscmd.c_str());
#endif
    }

    std::string d_to_str(double num) {
        char buffer[32];
        sprintf(buffer, "%g", num);
        std::string mret(buffer);
        return mret;
    }

    std::string i_to_str(int num) {
        char buffer[32];
        sprintf(buffer, "%d", num);
        std::string mret(buffer);
        return mret;
    }

    std::string col_to_hex(ChColor mcolor) {
        std::string mret;
        mret += "rgb \"#";
        std::stringstream mstreamR;
        mstreamR << std::setfill('0') << std::setw(2) << std::hex << (int)(mcolor.R * 255);
        std::string msr(mstreamR.str());
        mret += msr;
        std::stringstream mstreamG;
        mstreamG << std::setfill('0') << std::setw(2) << std::hex << (int)(mcolor.G * 255);
        std::string msg(mstreamG.str());
        mret += msg;
        std::stringstream mstreamB;
        mstreamB << std::setfill('0') << std::setw(2) << std::hex << (int)(mcolor.B * 255);
        std::string msb(mstreamB.str());
        mret += msb;
        mret += "\"";
        return mret;
    }

    std::string gpl_filename;
    std::string commandfile;
    std::vector<ChGnuPlotDataplot> plots;

    bool persist;
};

/// @} postprocess_module

}  // end namespace postprocess
}  // end namespace chrono

#endif
